import os
import shutil
import argparse
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter


class RAGPipeline:
    """
    一个完整的、高效的检索增强生成(RAG)流程封装。
    实现了索引的创建、加载和查询功能。
    """

    def __init__(self, config):
        """
        初始化 RAG 流程。
        在初始化时，会加载嵌入模型和本地的FAISS索引，以便后续快速查询。
        这避免了每次查询都重新加载的巨大开销。
        """
        self.config = config
        self.embeddings = self._load_embeddings()
        self.vector_store = self._load_vector_store()

    def _load_embeddings(self):
        """
        加载嵌入模型。
        这是整个流程中资源消耗较大的部分之一，只应在启动时执行一次。
        """
        print(f"--- 正在加载嵌入模型: {self.config['embedding_model']} ---")
        # HuggingFaceBgeEmbeddings 默认使用 CUDA (如果可用)
        model_kwargs = {'device': 'cuda'}
        encode_kwargs = {'normalize_embeddings': True}
        return HuggingFaceBgeEmbeddings(
            model_name=self.config['embedding_model'],
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs
        )

    def _load_vector_store(self):
        """
        从本地磁盘加载FAISS向量数据库。
        如果索引不存在，则返回 None。
        """
        index_path = self.config['faiss_index_dir']
        if not os.path.exists(index_path):
            print(f"警告：向量数据库目录 '{index_path}' 不存在。需要先创建索引。")
            return None

        print(f"--- 正在从 '{index_path}' 加载向量数据库 ---")
        try:
            return FAISS.load_local(
                index_path,
                self.embeddings,
                allow_dangerous_deserialization=True  # FAISS索引是pickle格式，需要允许
            )
        except Exception as e:
            print(f"加载向量数据库时出错: {e}")
            return None

    def create_index(self, force_recreate=False):
        """
        从知识库目录加载文档，创建并保存FAISS向量数据库。
        """
        index_path = self.config['faiss_index_dir']

        # 健壮性改进：如果索引已存在，除非强制重建，否则不执行
        if os.path.exists(index_path) and not force_recreate:
            print(f"向量数据库 '{index_path}' 已存在。如需重建，请使用 --rebuild 标志。")
            return

        if os.path.exists(index_path):
            print(f"警告：将要删除并重建已存在的索引目录 '{index_path}'...")
            shutil.rmtree(index_path)

        # 文档加载改进：使用 DirectoryLoader 加载所有 .txt 文件
        print(f"--- 开始从 '{self.config['knowledge_base_dir']}' 加载文档 ---")
        loader = DirectoryLoader(self.config['knowledge_base_dir'], glob="**/*.txt", show_progress=True)
        documents = loader.load()

        if not documents:
            print("错误：在知识库目录中没有找到任何 .txt 文件。")
            return

        print(f"成功加载 {len(documents)} 篇文档。")

        # 文本分割改进：使用 RecursiveCharacterTextSplitter
        # 它会尝试按段落、句子等更有意义的边界来分割文本，效果优于简单的字符分割。
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.config['chunk_size'],
            chunk_overlap=self.config['chunk_overlap']
        )
        docs = text_splitter.split_documents(documents)
        print(f"文档被分割成 {len(docs)} 个片段。")

        # 创建并构建 FAISS 向量数据库
        print("--- 正在创建 FAISS 向量数据库... ---")
        new_vector_store = FAISS.from_documents(docs, self.embeddings)

        # 保存到本地
        new_vector_store.save_local(index_path)
        self.vector_store = new_vector_store  # 更新当前实例的向量数据库
        print(f"--- 向量数据库成功创建并保存于 '{index_path}' 文件夹中！ ---")

    def search(self, query: str, k: int = 4):
        """
        执行相似度搜索。
        直接使用初始化时加载好的向量数据库，效率极高。
        """
        if not self.vector_store:
            print("错误：向量数据库未加载。请先创建或确保路径正确。")
            return []

        print("\n--- 开始执行相似度搜索 ---")
        results = self.vector_store.similarity_search(query, k=k)

        # 显示结果
        print(f"\n对于查询: '{query}', 找到以下 {len(results)} 个最相关的结果：")
        print("-" * 50)
        for i, doc in enumerate(results):
            print(f"结果 {i + 1}:")
            print(f"  内容: {doc.page_content}")
            print(f"  来源: {doc.metadata.get('source', 'N/A')}")
            print("-" * 20)

        return results


if __name__ == "__main__":
    # 使用 argparse 来处理命令行参数，更灵活
    parser = argparse.ArgumentParser(description="RAG Pipeline for Sichuan Travel Bot")
    parser.add_argument("--rebuild", action="store_true", help="如果指定，将强制重建向量索引。")
    parser.add_argument("--query", type=str, help="输入一个问题进行查询。")
    args = parser.parse_args()

    # --- 统一配置 ---
    config = {
        "knowledge_base_dir": "data_sample",
        "faiss_index_dir": "faiss_index",
        "embedding_model": "BAAI/bge-small-zh",
        "chunk_size": 400,
        "chunk_overlap": 50,
    }

    # 1. 初始化流程
    rag_pipeline = RAGPipeline(config)

    # 2. 如果需要，重建索引
    if args.rebuild:
        rag_pipeline.create_index(force_recreate=True)

    # 3. 如果有查询请求，执行查询
    if args.query:
        rag_pipeline.search(args.query)

    # 4. 如果没有特定指令，可以启动一个交互式查询循环
    if not args.rebuild and not args.query:
        if rag_pipeline.vector_store is None:
            print("请先使用 --rebuild 标志创建索引后再进行查询。")
        else:
            print("\n进入交互式查询模式。输入 '退出' 来结束程序。")
            while True:
                user_query = input("请输入您的问题: > ")
                if user_query.lower() in ["退出", "exit", "quit"]:
                    break
                rag_pipeline.search(user_query)